package com.bigdata.udf;

import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDTF;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.StructObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;

import java.util.ArrayList;
import java.util.List;

/**
 * 创建一个UDAF函数，
 * 接收2个字符串参数，第一个值是目标字符串，第二个是分隔符，用分割符分割第一个字符串，
 * 把分割的每个子串作为一行返回。
 */
public class MyUDTF extends GenericUDTF {


    @Override
    public StructObjectInspector initialize(StructObjectInspector argOIs) throws UDFArgumentException {
        List<String> structFieldNames = new ArrayList<>();
        structFieldNames.add("sub_str");
        List<ObjectInspector> objectInspectors = new ArrayList<>();
        objectInspectors.add(PrimitiveObjectInspectorFactory.javaStringObjectInspector);
        // UDTF 函数可以回返回多列，在return函数中传递的是List
        return ObjectInspectorFactory.getStandardStructObjectInspector(structFieldNames, objectInspectors);
    }

    // 当前的函数只返回一列，数组长度设置为1
    Object[] forwardListObj = new Object[1];

    @Override
    public void process(Object[] args) throws HiveException {
        String targetStr = args[0].toString();
        String splitStr = args[1].toString();

        // 返回的多行，分割的每个子串作为一行返回
        String[] splitArr = targetStr.split(splitStr);

        for(String str : splitArr) {
            // 设置一行数据中的所有列的值
            forwardListObj[0] = str;
            // forward 是写出一行数据，这一行数据可能有多列
            forward(forwardListObj);
        }
    }

    @Override
    public void close() throws HiveException {

    }
}
